import os
import re
import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import threading
import time
from openai import OpenAI

# AI client setup
client = OpenAI(
    api_key="0f56e8eb-8815-482d-8099-c5bef02f129a",
    base_url="https://ark.cn-beijing.volces.com/api/v3",
)


class BibTexCompleter:
    def __init__(self, root):
        self.root = root
        self.root.title("BibTeX Entry Completer")
        self.root.geometry("1100x750")

        # Set a nicer theme if available
        try:
            self.root.tk.call("source", "azure.tcl")
            self.root.tk.call("set_theme", "light")
        except:
            pass

        self.input_file_path = tk.StringVar()
        self.output_file_path = tk.StringVar()
        self.current_entry = tk.StringVar()
        self.progress_var = tk.DoubleVar()
        self.total_entries = tk.StringVar(value="0")
        self.completed_entries = tk.StringVar(value="0")
        self.processing = False
        self.cancel_requested = False

        self.setup_ui()

    def setup_ui(self):
        # Create styles
        style = ttk.Style()
        style.configure("Bold.TLabel", font=("", 10, "bold"))

        # Main container with padding
        main_container = ttk.Frame(self.root, padding="10")
        main_container.pack(fill=tk.BOTH, expand=True)

        # Title and description
        header_frame = ttk.Frame(main_container)
        header_frame.pack(fill=tk.X, pady=(0, 15))

        ttk.Label(header_frame, text="BibTeX Entry Completer", font=("", 16, "bold")).pack(anchor=tk.W)
        ttk.Label(header_frame,
                  text="This tool helps complete missing fields in BibTeX entries by querying an AI service.",
                  wraplength=1000).pack(anchor=tk.W, pady=(5, 0))

        # File selection section
        file_frame = ttk.LabelFrame(main_container, text="File Selection", padding="10")
        file_frame.pack(fill=tk.X, pady=(0, 10))

        # Input file row
        input_frame = ttk.Frame(file_frame)
        input_frame.pack(fill=tk.X, pady=5)
        ttk.Label(input_frame, text="Input BibTeX File:", width=15).pack(side=tk.LEFT)
        ttk.Entry(input_frame, textvariable=self.input_file_path, width=80).pack(side=tk.LEFT, fill=tk.X, expand=True,
                                                                                 padx=5)
        ttk.Button(input_frame, text="Browse...", command=self.browse_input_file).pack(side=tk.LEFT)

        # Output file row
        output_frame = ttk.Frame(file_frame)
        output_frame.pack(fill=tk.X, pady=5)
        ttk.Label(output_frame, text="Output BibTeX File:", width=15).pack(side=tk.LEFT)
        ttk.Entry(output_frame, textvariable=self.output_file_path, width=80).pack(side=tk.LEFT, fill=tk.X, expand=True,
                                                                                   padx=5)
        ttk.Button(output_frame, text="Browse...", command=self.browse_output_file).pack(side=tk.LEFT)

        # Buttons frame
        button_frame = ttk.Frame(file_frame)
        button_frame.pack(fill=tk.X, pady=(10, 0))

        self.process_button = ttk.Button(button_frame, text="Process BibTeX File", command=self.process_file,
                                         style="Accent.TButton")
        self.process_button.pack(side=tk.LEFT, padx=(0, 10))

        self.cancel_button = ttk.Button(button_frame, text="Cancel", command=self.cancel_processing, state=tk.DISABLED)
        self.cancel_button.pack(side=tk.LEFT)

        # Progress section
        progress_frame = ttk.LabelFrame(main_container, text="Progress", padding="10")
        progress_frame.pack(fill=tk.X, pady=(0, 10))

        # Current entry info
        current_frame = ttk.Frame(progress_frame)
        current_frame.pack(fill=tk.X, pady=(0, 10))

        ttk.Label(current_frame, text="Current Entry:", style="Bold.TLabel").pack(side=tk.LEFT, padx=(0, 5))
        ttk.Label(current_frame, textvariable=self.current_entry).pack(side=tk.LEFT)

        # Completion statistics
        stats_frame = ttk.Frame(progress_frame)
        stats_frame.pack(fill=tk.X, pady=(0, 10))

        ttk.Label(stats_frame, text="Completed:", style="Bold.TLabel").pack(side=tk.LEFT, padx=(0, 5))
        ttk.Label(stats_frame, textvariable=self.completed_entries).pack(side=tk.LEFT)
        ttk.Label(stats_frame, text=" of ").pack(side=tk.LEFT)
        ttk.Label(stats_frame, textvariable=self.total_entries).pack(side=tk.LEFT)
        ttk.Label(stats_frame, text=" entries").pack(side=tk.LEFT)

        # Progress bar
        self.progress_bar = ttk.Progressbar(progress_frame, variable=self.progress_var, length=1080, mode="determinate")
        self.progress_bar.pack(fill=tk.X)

        # Create a PanedWindow for log and preview
        paned = ttk.PanedWindow(main_container, orient=tk.HORIZONTAL)
        paned.pack(fill=tk.BOTH, expand=True)

        # Log section
        log_frame = ttk.LabelFrame(paned, text="Processing Log", padding="10")
        paned.add(log_frame, weight=1)

        # Create a frame for the text and scrollbar
        text_frame = ttk.Frame(log_frame)
        text_frame.pack(fill=tk.BOTH, expand=True)

        # Create scrollbar first
        log_scrollbar = ttk.Scrollbar(text_frame)
        log_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        # Create text widget with scrollbar
        self.log_text = tk.Text(text_frame, height=15, width=60, wrap=tk.WORD, yscrollcommand=log_scrollbar.set)
        self.log_text.pack(fill=tk.BOTH, expand=True)
        log_scrollbar.config(command=self.log_text.yview)

        # Add some initial instructions to the log
        self.log_text.insert(tk.END, "Instructions:\n")
        self.log_text.insert(tk.END, "1. Select an input BibTeX file using the 'Browse...' button\n")
        self.log_text.insert(tk.END, "2. Choose where to save the completed file\n")
        self.log_text.insert(tk.END, "3. Click 'Process BibTeX File' to start\n")
        self.log_text.insert(tk.END, "4. The log will show progress as each entry is processed\n\n")
        self.log_text.insert(tk.END, "Ready to start.\n")

        # Preview section
        preview_frame = ttk.LabelFrame(paned, text="Entry Preview", padding="10")
        paned.add(preview_frame, weight=1)

        preview_notebook = ttk.Notebook(preview_frame)
        preview_notebook.pack(fill=tk.BOTH, expand=True)

        # Original entry tab
        original_frame = ttk.Frame(preview_notebook, padding="5")
        preview_notebook.add(original_frame, text="Original Entry")

        # Create scrollbar for original
        original_scrollbar = ttk.Scrollbar(original_frame)
        original_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        # Create text widget for original entry
        self.original_text = tk.Text(original_frame, height=15, width=60, wrap=tk.WORD,
                                     yscrollcommand=original_scrollbar.set)
        self.original_text.pack(fill=tk.BOTH, expand=True)
        original_scrollbar.config(command=self.original_text.yview)

        # Completed entry tab
        completed_frame = ttk.Frame(preview_notebook, padding="5")
        preview_notebook.add(completed_frame, text="Completed Entry")

        # Create scrollbar for completed
        completed_scrollbar = ttk.Scrollbar(completed_frame)
        completed_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        # Create text widget for completed entry
        self.completed_text = tk.Text(completed_frame, height=15, width=60, wrap=tk.WORD,
                                      yscrollcommand=completed_scrollbar.set)
        self.completed_text.pack(fill=tk.BOTH, expand=True)
        completed_scrollbar.config(command=self.completed_text.yview)

        # Differences tab
        diff_frame = ttk.Frame(preview_notebook, padding="5")
        preview_notebook.add(diff_frame, text="Changes")

        # Create scrollbar for differences
        diff_scrollbar = ttk.Scrollbar(diff_frame)
        diff_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)

        # Create text widget for differences
        self.diff_text = tk.Text(diff_frame, height=15, width=60, wrap=tk.WORD, yscrollcommand=diff_scrollbar.set)
        self.diff_text.pack(fill=tk.BOTH, expand=True)
        diff_scrollbar.config(command=self.diff_text.yview)

        # Configure text tags for highlighting
        self.diff_text.tag_configure("added", foreground="green", font=("", 10, "bold"))
        self.diff_text.tag_configure("removed", foreground="red", font=("", 10, "bold"))
        self.diff_text.tag_configure("unchanged", foreground="black")

        # Status bar
        status_frame = ttk.Frame(self.root)
        status_frame.pack(fill=tk.X, side=tk.BOTTOM, pady=(5, 0))
        self.status_label = ttk.Label(status_frame, text="Ready", relief=tk.SUNKEN, anchor=tk.W, padding=(5, 2))
        self.status_label.pack(fill=tk.X)

    def browse_input_file(self):
        filename = filedialog.askopenfilename(
            title="Select BibTeX File",
            filetypes=[("BibTeX files", "*.bib"), ("All files", "*.*")]
        )
        if filename:
            self.input_file_path.set(filename)
            # Set a default output filename
            default_output = os.path.splitext(filename)[0] + "_completed.bib"
            self.output_file_path.set(default_output)
            self.status_label.config(text=f"Selected input file: {filename}")

            # Try to count entries in the file to inform the user
            try:
                with open(filename, 'r', encoding='utf-8') as file:
                    content = file.read()
                entries = self.parse_bibtex(content)
                self.total_entries.set(str(len(entries)))
                self.log(f"Found {len(entries)} BibTeX entries in the selected file.")
            except Exception as e:
                self.log(f"Error reading file: {str(e)}")

    def browse_output_file(self):
        filename = filedialog.asksaveasfilename(
            title="Save Completed BibTeX File",
            defaultextension=".bib",
            filetypes=[("BibTeX files", "*.bib"), ("All files", "*.*")]
        )
        if filename:
            self.output_file_path.set(filename)
            self.status_label.config(text=f"Selected output file: {filename}")

    def log(self, message):
        self.log_text.insert(tk.END, message + "\n")
        self.log_text.see(tk.END)
        self.root.update_idletasks()

    def update_preview(self, original, completed):
        # Update original entry preview
        self.original_text.delete(1.0, tk.END)
        self.original_text.insert(tk.END, original)

        # Update completed entry preview
        self.completed_text.delete(1.0, tk.END)
        self.completed_text.insert(tk.END, completed)

        # Update differences display
        self.update_diff_view(original, completed)

        self.root.update_idletasks()

    def update_diff_view(self, original, completed):
        self.diff_text.delete(1.0, tk.END)

        # Simple line-by-line diff
        original_lines = original.strip().split('\n')
        completed_lines = completed.strip().split('\n')

        # Create a dictionary of original fields for easier comparison
        original_fields = {}
        for line in original_lines:
            # Skip the first and last lines (entry type and closing brace)
            if line.startswith('@') or line.strip() == '}':
                continue

            # Try to extract field name and value
            field_match = re.match(r'\s*(\w+)\s*=\s*{(.+?)},?$', line)
            if field_match:
                field_name = field_match.group(1).strip()
                field_value = field_match.group(2).strip()
                original_fields[field_name] = field_value

        # Now process completed entry to show differences
        in_field = False
        for line in completed_lines:
            # Output the entry type and closing brace as is
            if line.startswith('@') or line.strip() == '}':
                self.diff_text.insert(tk.END, line + '\n', "unchanged")
                continue

            # Try to extract field name and value
            field_match = re.match(r'\s*(\w+)\s*=\s*{(.+?)},?$', line)
            if field_match:
                field_name = field_match.group(1).strip()
                field_value = field_match.group(2).strip()

                # Check if this field exists in the original entry
                if field_name in original_fields:
                    # Field exists in original - check if value changed
                    if original_fields[field_name] == field_value:
                        # Unchanged
                        self.diff_text.insert(tk.END, line + '\n', "unchanged")
                    else:
                        # Value changed
                        self.diff_text.insert(tk.END, line + '  // MODIFIED\n', "added")
                else:
                    # New field added
                    self.diff_text.insert(tk.END, line + '  // ADDED\n', "added")
            else:
                # Line doesn't match field pattern, output as is
                self.diff_text.insert(tk.END, line + '\n', "unchanged")

        # Check for fields in original that are missing in completed
        for field_name in original_fields:
            # Try to find the field in the completed entry
            if not re.search(rf'\s*{field_name}\s*=\s*{{', completed):
                self.diff_text.insert(tk.END, f"// REMOVED: {field_name} = {{{original_fields[field_name]}}}\n",
                                      "removed")

    def process_file(self):
        if self.processing:
            messagebox.showinfo("Info", "Processing is already in progress.")
            return

        input_path = self.input_file_path.get()
        output_path = self.output_file_path.get()

        if not input_path:
            messagebox.showerror("Error", "Please select an input BibTeX file.")
            return

        if not output_path:
            messagebox.showerror("Error", "Please select an output file location.")
            return

        # Clear displays
        self.log_text.delete(1.0, tk.END)
        self.original_text.delete(1.0, tk.END)
        self.completed_text.delete(1.0, tk.END)
        self.diff_text.delete(1.0, tk.END)

        # Reset progress
        self.progress_var.set(0)
        self.completed_entries.set("0")
        self.cancel_requested = False

        # Start processing in a separate thread to avoid freezing the UI
        self.processing = True
        self.process_button.config(state=tk.DISABLED)
        self.cancel_button.config(state=tk.NORMAL)
        self.status_label.config(text="Processing...")

        threading.Thread(target=self.process_file_thread, args=(input_path, output_path), daemon=True).start()

    def cancel_processing(self):
        if not self.processing:
            return

        response = messagebox.askyesno("Cancel Processing",
                                       "Are you sure you want to cancel processing?\n\nProgress so far will be saved.")
        if response:
            self.cancel_requested = True
            self.log("Cancellation requested. Waiting for current entry to complete...")
            self.status_label.config(text="Cancelling...")

    def process_file_thread(self, input_path, output_path):
        try:
            self.log(f"Reading BibTeX file: {input_path}")

            with open(input_path, 'r', encoding='utf-8') as file:
                content = file.read()

            # Parse BibTeX entries
            entries = self.parse_bibtex(content)

            if not entries:
                self.log("No BibTeX entries found in the file.")
                self.finish_processing("No BibTeX entries found.")
                return

            self.total_entries.set(str(len(entries)))
            self.log(f"Found {len(entries)} BibTeX entries.")

            # Process each entry
            completed_entries = []
            for i, entry in enumerate(entries):
                if self.cancel_requested:
                    self.log("Processing cancelled by user.")
                    break

                # Update progress
                progress_pct = ((i) / len(entries)) * 100
                self.progress_var.set(progress_pct)
                self.completed_entries.set(str(i))

                key = self.extract_key(entry)
                self.current_entry.set(f"{key} ({i + 1}/{len(entries)})")

                self.log(f"\n=== Processing entry [{i + 1}/{len(entries)}]: {key} ===")
                self.log(f"Original entry:\n{entry}\n")

                # Update preview with original entry
                self.update_preview(entry, "Querying AI for complete information...")

                # Query AI for complete information
                prompt = self.create_prompt(entry)
                ai_response = self.query_ai(prompt)

                # Parse AI response and update entry
                completed_entry = self.update_entry(entry, ai_response)
                completed_entries.append(completed_entry)

                # Update preview with completed entry
                self.update_preview(entry, completed_entry)

                self.log(f"Completed entry:\n{completed_entry}\n")
                self.log(f"Progress: {progress_pct:.1f}%")

            # Save completed entries to output file
            with open(output_path, 'w', encoding='utf-8') as file:
                file.write("\n\n".join(completed_entries))

            self.completed_entries.set(str(len(completed_entries)))
            self.progress_var.set(100 if not self.cancel_requested else (len(completed_entries) / len(entries)) * 100)

            if self.cancel_requested:
                status_msg = f"Cancelled. Processed {len(completed_entries)} of {len(entries)} entries."
                self.finish_processing(status_msg)
                self.log(f"\nProcessing cancelled. Partial results saved to: {output_path}")
                messagebox.showinfo("Process Cancelled", f"{status_msg}\nPartial results saved to: {output_path}")
            else:
                self.current_entry.set("All entries processed")
                status_msg = f"Completed! Processed {len(entries)} entries."
                self.finish_processing(status_msg)
                self.log(f"\nAll entries processed successfully!")
                self.log(f"Output saved to: {output_path}")
                messagebox.showinfo("Success",
                                    f"BibTeX file processing completed.\nAll {len(entries)} entries processed.\nSaved to: {output_path}")

        except Exception as e:
            self.log(f"Error: {str(e)}")
            self.finish_processing(f"Error: {str(e)}")
            messagebox.showerror("Error", f"An error occurred: {str(e)}")

    def finish_processing(self, status_message):
        self.processing = False
        self.process_button.config(state=tk.NORMAL)
        self.cancel_button.config(state=tk.DISABLED)
        self.status_label.config(text=status_message)

    def parse_bibtex(self, content):
        # Improved BibTeX parser
        # First, try to match entries with standard formatting (from @ to closing })
        entries = []

        # Look for entries that start with @ and end with a line containing only }
        entry_pattern = r'(@\w+\{[^@]+?(?:\n\s*\}))'
        found_entries = re.findall(entry_pattern, content, re.DOTALL)
        entries.extend(found_entries)

        # If no entries found, try more flexible patterns
        if not entries:
            # Try a pattern that captures entries where the closing brace might be on the same line as the last field
            entry_pattern = r'(@\w+\{[^@]+?\})'
            found_entries = re.findall(entry_pattern, content, re.DOTALL)
            entries.extend(found_entries)

        # Clean up entries
        cleaned_entries = []
        for entry in entries:
            # Remove any trailing whitespace or commas before the closing brace
            cleaned_entry = re.sub(r',\s*\}$', '\n}', entry.strip())
            cleaned_entries.append(cleaned_entry)

        return cleaned_entries

    def extract_key(self, entry):
        # Extract the citation key from the entry
        key_match = re.search(r'@\w+\{([^,]+)', entry)
        if key_match:
            return key_match.group(1).strip()
        return "Unknown"

    def create_prompt(self, entry):
        # Create a detailed prompt for the AI with clear formatting instructions
        prompt = (
            "I need to complete a BibTeX entry with missing bibliographic information. Help me find all standard fields "
            "for this publication to make it as complete as possible.\n\n"
            "Here's the incomplete entry:\n\n"
            f"{entry}\n\n"
            "Please search for this publication and provide a complete BibTeX entry with all standard fields "
            "(journal, volume, number, pages, month, year, doi, url, publisher, address, etc.).\n\n"
            "Your response MUST use this exact format to enable automatic parsing:\n\n"
            "BIBTEX_ENTRY_START\n"
            "@type{key,\n"
            "    field1={value1},\n"
            "    field2={value2},\n"
            "    ...\n"
            "}\n"
            "BIBTEX_ENTRY_END\n\n"
            "Important requirements:\n"
            "1. Use the EXACT markers BIBTEX_ENTRY_START and BIBTEX_ENTRY_END\n"
            "2. Keep the original citation key and entry type unless clearly incorrect\n"
            "3. Use ONLY curly braces {} for field values, not quotes\n"
            "4. Include ALL standard BibTeX fields appropriate for this publication type\n"
            "5. If you cannot find exact information, make a reasonable guess based on similar publications\n"
            "6. Ensure proper formatting with each field on a new line and indented\n"
            "7. Check that all author names are properly formatted with correct 'and' separators\n"
            "8. Verify the entry type (article, inproceedings, etc.) is appropriate for the publication\n"
        )
        return prompt

    def query_ai(self, prompt):
        self.log("Querying AI for complete information...")
        attempts = 0
        max_attempts = 3

        while attempts < max_attempts:
            try:
                start_time = time.time()

                completion = client.chat.completions.create(
                    model="deepseek-r1-250120",
                    messages=[
                        {"role": "system",
                         "content": "You are a helpful assistant for academic researchers. You excel at finding complete bibliographic information for academic publications and formatting it correctly in BibTeX format."},
                        {"role": "user", "content": prompt},
                    ],
                )

                response = completion.choices[0].message.content
                elapsed_time = time.time() - start_time

                self.log(f"Received response from AI in {elapsed_time:.2f} seconds.")
                return response

            except Exception as e:
                attempts += 1
                self.log(f"Error querying AI (attempt {attempts}/{max_attempts}): {str(e)}")
                if attempts < max_attempts:
                    self.log("Retrying in 5 seconds...")
                    time.sleep(5)
                else:
                    self.log("Maximum retry attempts reached. Skipping this entry.")
                    return None

    def update_entry(self, original_entry, ai_response):
        if not ai_response:
            self.log("No valid response from AI. Keeping original entry.")
            return original_entry

        # Extract the BibTeX entry from the AI response using the markers
        bibtex_pattern = r'BIBTEX_ENTRY_START\s*\n(.*?)\s*BIBTEX_ENTRY_END'
        entry_match = re.search(bibtex_pattern, ai_response, re.DOTALL)

        if entry_match:
            completed_entry = entry_match.group(1).strip()
            self.log("Successfully extracted completed entry using the specified markers.")
            return completed_entry
        else:
            # Try alternative patterns in case the AI didn't follow the exact format
            self.log("Could not find the exact markers. Trying alternative parsing methods...")

            # Look for bibtex code blocks
            bibtex_pattern = r'```bibtex\s*\n(.*?)\s*```'
            entry_match = re.search(bibtex_pattern, ai_response, re.DOTALL)

            if entry_match:
                completed_entry = entry_match.group(1).strip()
                self.log("Extracted entry from bibtex code block.")
                return completed_entry
            else:
                # Try without the bibtex tag
                bibtex_pattern = r'```\s*\n(.*?)\s*```'
                entry_match = re.search(bibtex_pattern, ai_response, re.DOTALL)

                if entry_match:
                    completed_entry = entry_match.group(1).strip()
                    self.log("Extracted entry from generic code block.")
                    return completed_entry
                else:
                    # If no clear code formatting, try to extract @ entry directly
                    bibtex_pattern = r'(@\w+\{[^@]+\})'
                    entry_match = re.search(bibtex_pattern, ai_response, re.DOTALL)

                    if entry_match:
                        completed_entry = entry_match.group(1).strip()
                        self.log("Extracted entry directly from response using @ pattern.")
                        return completed_entry
                    else:
                        self.log("Could not parse AI response. Keeping original entry.")
                        self.log(f"AI response was:\n{ai_response}")
                        return original_entry


# Main application
def main():
    root = tk.Tk()
    app = BibTexCompleter(root)
    root.mainloop()


if __name__ == "__main__":
    main()